# -*- coding: utf-8 -*-
"""
Created on Thu Aug 15 02:10:36 2013

@author: John Edwardson
"""
import numpy as np
import ConfigParser as cp
from Converters import UnitConverter, ScaleConverter


class State(object):
    unitC = UnitConverter()
    unitC.buildConversionTable()

    def __init__(self, fileinfo, label, style, molecule):
        '''State is a class that holds all the curves and annotations corresponding to a given 
        electronic state.  '''
        self.fileinfo = fileinfo
        self.path = fileinfo[0]
        self.label = label
        self.style = style
        self.units = dict(Radius=fileinfo[1], Energy=fileinfo[2], Scale=fileinfo[3])
        self.name = molecule
        try:
            self.comment = fileinfo[4]
            self.skip = int(fileinfo[5])
        except IndexError:
            self.comment = '#'
            self.skip = 0
        
        self.data = np.loadtxt(self.path, comments=self.comment, skiprows=self.skip)
            
        self.color = self.style[0]
        self.ls = self.style[1]
        self.lw = self.style[2]        
        
        self.curves = []
        self.annotes = []
        self.parseLabel()
        
    def parseLabel(self):
        mlegend = {'S': 'Sigma', 'D': 'Delta', 'P': 'Pi', 'F': 'Phi'}        
        self.mult = self.label[0]
        self.term = mlegend[self.label[1]]
        self.rsym = self.label[-1]
        if '+' in self.label: self.sym = '+'
        elif '-' in self.label: self.sym = '-'
        else: self.sym = ''        
        
    def convert(self, outunits):
        '''Converts energy, radius, and scale to proper output units'''        
        #Get ground state energy and units from molecules file
        parser2 = cp.ConfigParser()
        parser2.read('molecules.cfg')
        gse = parser2.getfloat(self.name, "Ground State")
        gsin = parser2.get(self.name, 'Units')       
        #Bind input units locally for convenience
        rin = self.units['Radius']
        ein = self.units["Energy"]
        esin = self.units['Scale']        
        plotdata = self.data        
        if outunits: #if no output units are specified, then we need not do any conversion
            rout = outunits['Radius']
            eout = outunits['Energy']
            esout = outunits['Scale']
            gse *= State.unitC.conversion_factor('Energy', gsin, eout)
            sc = ScaleConverter(gse)
            plotdata[:,0] *= State.unitC.conversion_factor('Radius', rin, rout)
            plotdata[:,1:] *= State.unitC.conversion_factor('Energy', ein, eout) 
            plotdata[:,1:] += sc.conversion_offset(esin, esout)            
        return plotdata

    def plotPEC(self, ax, outunits = None): 
        '''Plots the potential energy curves associated with this state'''   
        plotdata = self.convert(outunits)
        for i in range(1,plotdata.shape[1]):
            additem, = ax.plot(plotdata[:,0], plotdata[:,i], color=self.color, ls=self.ls, lw=self.lw, picker=5.0)
            self.curves.append(additem)
        #initialize annotes
        self.annotes = [None for i in range(len(self.curves))]

    def OnOff(self, curve, flag):
        '''Sets visibility of a single curve'''
        idx = self.curves.index(curve)                
        self.curves[idx].set_visible(flag)
        if self.annotes[idx] != None:
            self.annotes[idx].set_visible(flag)
            
    def WaxOnWaxOff(self, flag):
        '''Sets visibility of all curves  of a given state'''
        for curve in self.curves:
            self.OnOff(curve, flag)
    
    def generateTex(self, i):
        return '$(' + str(i) + ')\,^{' + self.mult + '}\\'      \
                + self.term + '^{' + self.sym + '}' + '_{' + self.rsym + '}$'
                    
    def create_annote(self, ax, curve, xydata, tc):
        index = self.curves.index(curve)
        annlabel = self.generateTex(index+1)
        self.annotes[index] = ax.annotate(annlabel, xy=xydata, xytext = tc, ha='left', va='bottom', 
                                 fontsize=14, arrowprops=dict(arrowstyle='->'), picker = 5.0) 
    def removeAnnote(self, annote):
        idx = self.annotes.index(annote)
        self.annotes[idx].set_visible(False)
        self.annotes[idx] = None
        
    def get_visible(self):
        '''Returns visible kwd for all curves'''
        return [curve.get_visible() for curve in self.curves]

 
    



        


        

        
        
